package utils

import (
	"context"
	"github.com/cloudwego/eino-ext/components/model/ollama"
	"github.com/cloudwego/eino/components/prompt"
	"github.com/cloudwego/eino/schema"
)

type ChatSession struct {
	Role            string
	Style           string
	target          string
	Question        string
	UserSchema      []string
	AssistantSchema []string
	template        *prompt.DefaultChatTemplate
	result          *schema.Message
	err             error
}

func NewChat(role string, style string, target string, parameters ...any) *ChatSession {
	session := &ChatSession{
		Role:   role,
		Style:  style,
		target: target,
	}
	if len(parameters) > 0 {
		if askList := parameters[0]; askList != nil {
			session.UserSchema = askList.([]string)
		}
		if answerList := parameters[0]; answerList != nil {
			session.AssistantSchema = answerList.([]string)
		}
	}
	session.template = prompt.FromMessages(schema.FString,
		schema.SystemMessage("你是一个｛role｝。你需要用｛style｝的语气回答问题。你的目标是｛target｝。"),
		schema.MessagesPlaceholder("chat_history", true),
		schema.UserMessage("问题：｛question｝"),
	)
	return session
}

func (cs *ChatSession) InsertHistory(ask []string, answer []string) {
	for i, content := range ask {
		cs.UserSchema = append(cs.UserSchema, content)
		cs.AssistantSchema = append(cs.AssistantSchema, answer[i])
	}
}

func (cs *ChatSession) GetError() error {
	return cs.err
}

func (cs *ChatSession) GetResult() *schema.Message {
	return cs.result
}

func (cs *ChatSession) Send(question string) *ChatSession {
	var histories []*schema.Message
	if len(cs.UserSchema) > 0 {
		for i, content := range cs.UserSchema {
			histories = append(histories, &schema.Message{Role: schema.User, Content: content})
			histories = append(histories, &schema.Message{Role: schema.Assistant, Content: cs.AssistantSchema[i]})
		}
	}
	messages, err := cs.template.Format(context.Background(), map[string]any{
		"role":         cs.Role,
		"style":        cs.Style,
		"target":       cs.target,
		"question":     question,
		"chat_history": histories,
	})
	if err != nil {
		cs.err = err
		return cs
	}
	return cs.SendChatModelMessages(messages)
}

func (cs *ChatSession) SendChatModelMessages(messages []*schema.Message) *ChatSession {

	//chatModel, err := openai.NewChatModel(context.Background(), &openai.ChatModelConfig{
	//	BaseURL: "https://openrouter.ai/api/v1",
	//  Model:   "google/gemini-2.0-flash-thinking-exp-1219:free",
	//  APIKey:  os.Getenv("OPENAI_API_KEY"),
	//})
	chatModel, err := ollama.NewChatModel(context.Background(), &ollama.ChatModelConfig{
		BaseURL: "http://192.168.8.26:11434",
		Model:   "deepseek-r1:7b",
	})
	if err != nil {
		cs.err = err
		return cs
	}
	result, err := chatModel.Generate(context.Background(), messages)
	if err != nil {
		cs.err = err
		return cs
	}
	cs.result = result
	return cs
}
